from langchain.agents import Tool, initialize_agent, AgentType
from langchain.prompts import MessagesPlaceholder
from langchain.memory import ConversationBufferMemory
from langchain.chat_models import ChatOpenAI  # 或使用本地模型如LlamaCPP
from langchain.tools import DuckDuckGoSearchRun
from langchain.llms import LlamaCpp
import os
from dotenv import load_dotenv

# 加载环境变量（如OPENAI_API_KEY）
load_dotenv()

# ===== 1. 定义工具 =====
# 示例工具1：搜索互联网
search = DuckDuckGoSearchRun()

# 示例工具2：本地任务管理（模拟数据库）
task_db = {"tasks": []}

def add_task(task: str) -> str:
    """添加任务到数据库"""
    task_db["tasks"].append({"task": task, "status": "待办"})
    return f"已添加任务: {task}"

def list_tasks() -> str:
    """列出所有任务"""
    if not task_db["tasks"]:
        return "当前没有任务。"
    return "\n".join(
        f"{i+1}. [{t['status']}] {t['task']}"
        for i, t in enumerate(task_db["tasks"])
    )

# 将函数包装为LangChain工具
tools = [
    Tool(name="Search", func=search.run, description="搜索互联网信息"),
    Tool(name="AddTask", func=add_task, description="添加任务到待办列表"),
    Tool(name="ListTasks", func=list_tasks, description="列出所有任务"),
]

# ===== 2. 初始化智能体 =====
# 使用GPT-3.5（需配置API密钥）或本地模型
#llm = ChatOpenAI(temperature=0, model="gpt-3.5-turbo")
llm = LlamaCpp(model_path="/path/to/gguf-model.bin")

# 记忆模块（保存对话历史）
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

# 初始化智能体
agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.OPENAI_FUNCTIONS,  # 或ZERO_SHOT_REACT_DESCRIPTION
    memory=memory,
    verbose=True,  # 打印思考过程
)

# ===== 3. 运行智能体 =====
if __name__ == "__main__":
    print("🤖 任务助手已启动！输入'退出'结束对话。")
    while True:
        user_input = input("\n你: ")
        if user_input.lower() == "退出":
            break
        # 调用智能体处理输入
        response = agent.run(user_input)
        print(f"助手: {response}")